// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/executors/io_thread_pool_executor.h>

#include <pollux/exec/table_writer.h>
#include <pollux/exec/trace_util.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/tool/trace/table_writer_replayer.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::exec::test;

namespace kumo::pollux::tool::trace {
namespace {

std::shared_ptr<connector::hive::HiveInsertTableHandle>
makeHiveInsertTableHandle(
    const core::TableWriteNode* node,
    std::string targetDir) {
  const auto tracedHandle =
      std::dynamic_pointer_cast<connector::hive::HiveInsertTableHandle>(
          node->insertTableHandle()->connectorInsertTableHandle());
  const auto inputColumns = tracedHandle->inputColumns();
  const auto compressionKind =
      tracedHandle->compressionKind().value_or(common::CompressionKind_NONE);
  const auto storageFormat = tracedHandle->storageFormat();
  const auto serdeParameters = tracedHandle->serdeParameters();
  const auto writerOptions = tracedHandle->writerOptions();
  return std::make_shared<connector::hive::HiveInsertTableHandle>(
      inputColumns,
      std::make_shared<connector::hive::LocationHandle>(
          targetDir,
          targetDir,
          connector::hive::LocationHandle::TableType::kNew),
      storageFormat,
      tracedHandle->bucketProperty() == nullptr
          ? nullptr
          : std::make_shared<connector::hive::HiveBucketProperty>(
                *tracedHandle->bucketProperty()),
      compressionKind,
      std::unordered_map<std::string, std::string>{},
      writerOptions);
}

std::shared_ptr<core::InsertTableHandle> createInsertTableHanlde(
    const std::string& connectorId,
    const core::TableWriteNode* node,
    std::string targetDir) {
  return std::make_shared<core::InsertTableHandle>(
      connectorId, makeHiveInsertTableHandle(node, std::move(targetDir)));
}
} // namespace

core::PlanNodePtr TableWriterReplayer::createPlanNode(
    const core::PlanNode* node,
    const core::PlanNodeId& nodeId,
    const core::PlanNodePtr& source) const {
  const auto* tableWriterNode = dynamic_cast<const core::TableWriteNode*>(node);
  POLLUX_CHECK_NOT_NULL(tableWriterNode);
  const auto insertTableHandle =
      createInsertTableHanlde("test-hive", tableWriterNode, replayOutputDir_);
  return std::make_shared<core::TableWriteNode>(
      nodeId,
      tableWriterNode->columns(),
      tableWriterNode->columnNames(),
      tableWriterNode->aggregationNode(),
      insertTableHandle,
      tableWriterNode->hasPartitioningScheme(),
      TableWriteTraits::outputType(tableWriterNode->aggregationNode()),
      tableWriterNode->commitStrategy(),
      source);
}
} // namespace kumo::pollux::tool::trace
